Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training
MoE 训练中优化器状态该放在哪里？分层状态分配实现内存高效训练
🧭 MoE 训练内存优化
Independent Researcher
针对 MoE 三种参数群体，差异化分配优化器状态，使内存降至原来的 2.6%，且困惑度更优
📌 工作参考价值（4 点）
- 大规模 MoE 推荐模型训练：在训练广告/内容推荐的 MoE 模型时，专家层参数通常占 95% 以上，可仅用 factored second moment（无 momentum），骨干网络保留动量，大幅降低显存，使单卡训练更大 batch 或更多专家成为可能。
- 骨干与专家差异化优化：业务中可借鉴“骨干保留动量、专家不保留动量”的思路，骨干提取通用特征需要稳定收敛，而大量专家可通过轻量状态快速适应不同领域/任务。
- 路由器仅需精确二阶矩：推荐 MoE 中路由器通常极小，保持完整 Adam 状态几乎无开销，且对负载均衡有利（论文中负载均衡接近均匀下限）。
- 从显存到精度无 trade-off：实验表明分层状态分配不损害精度，甚至更好，可直接迁移到现有 MoE 训练流程，无需担心性能回退。